import baostock as bs
import pandas as pd
import requests
import time, sched
import datetime
import concurrent.futures as futures
from multiprocessing import Pool
import baostock.common.context as conx

def login():
    # 登陆系统
    lg = bs.login()
    # 显示登陆返回信息
    print('login respond error_code:' + lg.error_code)
    print('login respond  error_msg:' + lg.error_msg)
    return lg


def get_stocks():
    return bs.query_stock_industry()


def logout():
    bs.logout()


def query_history_k_data_plus(code, start_date, end_date):
    rs = bs.query_history_k_data_plus(code,
                                      "date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,"
                                      "tradestatus,pctChg,isST",
                                      start_date=start_date, end_date=end_date,
                                      frequency="d", adjustflag="3")
    data_list = []
    while (rs.error_code == '0') & rs.next():
        # 获取一条记录，将记录合并在一起
        data_list.append(rs.get_row_data())
    data = pd.DataFrame(data_list, columns=rs.fields)
    return data



def _spider(ev_time,
            scheduler,
            spider_stocks,
            stocks_history,
            alert_condition,
            alert_method,
            executor
            ):
    scheduler.enter(ev_time, 0, _spider, (ev_time, scheduler, spider_stocks, stocks_history, alert_condition, alert_method, executor))
    _results = []
    for i in range(0, len(spider_stocks)):
        history = stocks_history[i]
        stock = spider_stocks[i]
        _results.append(executor.submit(alert_condition, history, stock))
    for _result in _results:
        res, code = _result.result()
        alert_method(res, code)


def spider(spider_stocks=None,
           alert_method=None,
           alert_condition=None,
           history_operator=None,
           scheduler=sched.scheduler(time.time, time.sleep),
           for_time=10
           ):
    if spider_stocks is None \
            or alert_condition is None:
        return
    if alert_method is None:
        def _print(res, stock):
            if not res:
                return
            if res == 1:
                print('请买入：', stock)
            if res == 2:
                print('请卖出:', stock)
        alert_method = _print
    if history_operator is None:
        def default_history(stock):
            today = datetime.datetime.today()
            last_week = today + datetime.timedelta(-30)
            return query_history_k_data_plus(stock,
                                                last_week.strftime('%Y-%m-%d'),
                                                today.strftime('%Y-%m-%d')
                                                )
        history_operator = default_history
    login()
    stocks_history = [history_operator(stock)
                      for stock in spider_stocks]
    logout()
    # 创建线程，然后对操作
    with futures.ThreadPoolExecutor(len(spider_stocks)) as executor:
        _spider(for_time, scheduler,spider_stocks, stocks_history,alert_condition, alert_method, executor)
        scheduler.run()


class GetLines(object):

    def __init__(self):
        self._lg = None

    def __enter__(self):
        self._lg = login()

    def __exit__(self, exc_type, exc_val, exc_tb):
        logout()

    def get_kl(self, code, start, end):

        if self._lg is None or self._lg.error_code == '0':
            return query_history_k_data_plus(code, start, end)
        else:
            print('登陆失败、请重试')
            return None


if __name__ == '__main__':
    pool = Pool(4)
    resutls = []
    for c in ['sh.601015','sh.601016','sh.601018','sh.601019']:
        resutls.append(pool\
                       .apply_async(query_history_k_data_plus,
                                        (c, '2019-12-01', '2019-12-23', user, ), ))
    pool.close()
    pool.join()
    for result in resutls:
        print(result.get())
